"""
@author: 景云鹏
@email: 310491287@qq.com
@date: 2022/4/28
"""

import random
from typing import NamedTuple, Sequence

from numpy import *
from matplotlib import pyplot as plt

from mpl_toolkits.mplot3d import Axes3D


class Interval(NamedTuple):
    lower_bound: float
    upper_bound: float


class Domain:
    def __init__(self, intervals):
        self.intervals = intervals

    def check(self, x):
        if len(self.intervals) == 1 and len(x) >= 2:
            self.intervals *= len(x)
        for i, xi in enumerate(x):
            if xi < self.intervals[i * 2] or xi > self.intervals[i * 2 + 1]:
                return False
        return True


class Function:
    functions = set()

    def __init__(self, dim=None, min_value=0, intervals=None):
        self.f = None
        self.name = ''

        self.min_value = min_value
        self.intervals = intervals or [-32.768, 32.768]

        self.allow_set_dim = True
        self.dim = dim or 2
        self.allow_set_dim = dim is None

    def __setattr__(self, k, v):
        if k == 'dim':
            if not self.allow_set_dim:
                raise ValueError('方法{}不允许修改维度'.format(self.f.name))
            if not isinstance(self.intervals[0], Sequence):
                self.intervals = [self.intervals] * v
            if len(self.intervals) != v:
                self.intervals = [self.intervals[0]] * v
            if self.name:
                self.name = '{}({})'.format(
                    self.f.__name__, v
                )

        return super().__setattr__(k, v)

    def __call__(self, arg):
        if callable(arg):
            self.f = arg
            self.name = '{}({})'.format(
                arg.__name__, self.dim
            )
            self.functions.add(self)
            return self
        return self.f(arg)

    def random_x(self):
        x = []
        for interval in self.intervals:
            x.append(random.uniform(interval[0], interval[1]))
        return asarray(x)

    def check_bounds(self, x):
        for i, (lb, ub) in enumerate(self.intervals):
            x[i] = max((x[i], lb))
            x[i] = min((x[i], ub))

    def figure(self, points=None):
        #assert self.dim == 2, '仅支持2维函数绘制3维图像'
        b1, b2 = self.intervals[:2]
        x1 = arange(b1[0], b1[1], 0.2)
        x2 = arange(b2[0], b2[1], 0.2)
        x1, x2 = meshgrid(x1, x2)
        z = self.f((x1, x2))

        figure = plt.figure()
        figure.suptitle(self.name)
        axes = Axes3D(figure, auto_add_to_figure=False)
        figure.add_axes(axes)
        if points:
            for point, best in points:
                plt.cla()
                axes.plot_surface(x1, x2, z, alpha=0.1, cstride=1, rstride=1, cmap='rainbow')
                axes.scatter(*point, c='black')
                axes.scatter(*best, c='red', marker='*')
                plt.pause(0.1)
        else:
            axes.plot_surface(x1, x2, z, cstride=1, rstride=1, cmap='rainbow')
        plt.show()
